# -*- encoding: utf-8 -*-
'''
@File    :   cwru_reader.py
@Time    :   2022/3/23 8:56
@Author  :   ZhangChaoYang
@Desc    :   把CWRU原始数据（只使用驱动端的数据）文件，经过各种转换（FFT、CWT等）变成numpy ndarray，存入文件。美国凯斯西储大学轴承数据 https://engineering.case.edu/bearingdatacenter/download-data-file
'''

import os, sys
import random

sys.path.insert(0, os.getcwd())
from data_process import mat_reader
from data_process.mat_reader import convertDir1d, convertDir2d, take_a_look_1d, take_a_look_2d
import re
from absl import flags
from absl import app

FLAGS = flags.FLAGS

flags.DEFINE_string('source_data', "E:\zcy\数据集\CWRU", '原始数据存在目录')

fs = 12e3  # 采样频率12kHz
WC = ["Normal", "Ball", "InnerRace", "OuterRace6"]
regex1 = re.compile("^([a-zA-Z]+)_\d+\.mat")
regex2 = re.compile("0\.\d+-([a-zA-Z]+\d?)\.mat")


class CWRUReader(mat_reader.MatFileReader):
    def __init__(self, infile, trans1d="", trans2d="", work_load=""):
        super(CWRUReader, self).__init__(infile, "cwru", trans1d, trans2d, work_load)

    def get_label(self):
        basename = self.get_basename()
        if basename.lower().count("normal") > 0:  # 文件名中包含normal是正常样本
            return 1
        else:
            return 0

    def window_size(self):
        return 1024

    def read1d(self):
        if self.work_load:
            basename = self.get_basename()
            wc = regex1.findall(basename) or regex2.findall(basename)
            if len(wc) == 0 or wc[0] != self.work_load:
                return None
        dic = self.read_file_to_dict()

        # print(dic.keys())#['__header__', '__version__', '__globals__', 'X097_DE_time', 'X097_FE_time', 'X097RPM']
        for key in dic.keys():
            if key.count("_DE_") > 0:  # 只使用驱动端的数据
                array = dic[key][:, 0]
                x = self.trans1d(array, 1, fs)
                return x

    def read2d(self):
        if self.work_load:
            basename = self.get_basename()
            wc = regex1.findall(basename) or regex2.findall(basename)
            if len(wc) == 0 or wc[0] != self.work_load:
                return None
        dic = self.read_file_to_dict()
        if random.random() < 0.5:
            return None
        for key in dic.keys():
            if key.count("_DE_") > 0:  # 只使用驱动端的数据
                array = dic[key][:, 0]
                x = self.trans2d(array, 1, fs)
                return x


def main(argv):
    for trans1d in ['original', 'stat', 'fft']:
        reader = CWRUReader("", trans1d=trans1d)
        convertDir1d(reader, [FLAGS.source_data])
        for wc in WC:
            reader = CWRUReader("", trans1d=trans1d, work_load=wc)
            convertDir1d(reader, [FLAGS.source_data])
        take_a_look_1d(reader.name, trans1d)

    for trans2d in ['sfft', 'cwt']:
        reader = CWRUReader(FLAGS.source_data, trans2d=trans2d)
        convertDir2d(reader, [FLAGS.source_data])
        for wc in WC:
            reader = CWRUReader("", trans2d=trans2d, work_load=wc)
            convertDir2d(reader, [FLAGS.source_data])
        take_a_look_2d(reader.name, trans2d)


if __name__ == '__main__':
    app.run(main)

# python .\data_process\cwru_reader.py --source_data E:\zcy\数据集\CWRU
